#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-i', '--input', help='Input file', default='mz.360x180x26.h0.nc')
parser.add_argument('-t', '--time-step', type=int, help='Time step to plot', default=-1)
parser.add_argument('-l', '--loop', help='Loop plot', action='store_true')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

fig = plt.figure(figsize=(10, 6))
proj = ccrs.PlateCarree(central_longitude=180)

def plot(data_path, fig):
	f = addfile(data_path).isel(time=args.time_step)

	p = f.ph

	time = parse_datetime(f.time)
	fig.suptitle(time.strftime('%Y-%m-%d'), fontsize=8)

	i = 0

	i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
	u700 = vinterp_p(p, f.u, 700 * units.hPa)
	plot_contour_map(ax, u700[0,:,:], cmap='jet', levels=np.linspace(-15, 50, 14), left_string='U at 700 hPa')

	i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
	v700 = vinterp_p(p, f.v, 700 * units.hPa)
	plot_contour_map(ax, v700[0,:,:], cmap='jet', left_string='V at 700 hPa', levels=np.linspace(-35, 20, 12))

	i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
	z700 = vinterp_p(p, f.gz, 700 * units.hPa) / 9.80616
	plot_contour_map(ax, z700[0,:,:], cmap='jet', left_string='Z at 700 hPa', levels=np.linspace(2400, 3400, 11))

	i += 1; ax = fig.add_subplot(2, 2, i, projection=proj)
	t700 = vinterp_p(p, f.t, 700 * units.hPa)
	plot_contour_map(ax, t700[0,:,:], cmap='jet', left_string='T at 700 hPa', levels=np.linspace(270, 303, 12))

	f.close()

if args.loop:
	loop_plot(args.input, fig, plot)
else:
	plot(args.input, fig)
	plt.tight_layout()
	if args.output:
	  plt.savefig(args.output, dpi=300)
	else:
	  plt.show()
